package cn.kivensoft.util;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

/** jwt处理辅助函数
 * @author kiven
 * @version 1.0
 */
final public class JwtHelper {
	public static final int ERR_TOKEN_IS_NULL = 1; // token为空
	public static final int ERR_BAD_TOKEN     = 2; // token格式不正确
	public static final int ERR_NOT_HS256     = 3; // token加密类型不是HS256
	public static final int ERR_NOT_JWT       = 4; // token类型不是jwt
	public static final int ERR_BAD_ISSUER    = 5; // token发行者错误
	public static final int ERR_EXPIRE        = 6; // token过期
	public static final int ERR_BAD_SIGN      = 7; // token签名错误

	private static final String HS256 = "HS256";
	private static final String JWT = "JWT";
	private static final String HMAC_SHA256 = "HmacSHA256";
	private static final String ISSUER_KEY = "iss"; // 签发者
	private static final String EXP_KEY = "exp"; // 过期时间(秒为单位)

	// {"alg":"HS256","typ":"JWT"}
	private static final String JWT_TYPE_BASE64 = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9";

	private JwtHelper() {}

	public static String encode(String json, String key) throws Exception {
		Fmt f = Fmt.get().append(JWT_TYPE_BASE64).append('.')
				.appendBase64Url(json.getBytes(StandardCharsets.UTF_8), false);
		byte[] encData = f.toBytes();
		try {
			byte[] enc = encrypt(Strings.toBytes(key), encData);
			return f.append('.').appendBase64Url(enc, false).toString();
		} finally {
			f.recycle();
		}
	}

	/** 编码data为jwt token令牌
	 * @param data 编码内容
	 * @param issuer 发布者
	 * @param key 密钥
	 * @param exp 截止过期时间
	 * @return 生成的令牌
	 */
	public static String encode(Map<String, Object> data, String issuer, String key, Date exp) throws Exception {
		HashMap<String, Object> claim = new HashMap<>(data);
		if (issuer != null && !issuer.isEmpty())
			claim.put(ISSUER_KEY, issuer);
		claim.put(EXP_KEY, exp.getTime() / 1_000L);
		String json = Fmt.toJson(claim);
		Fmt f = Fmt.get().append(JWT_TYPE_BASE64).append('.')
				.appendBase64Url(json.getBytes(StandardCharsets.UTF_8), false);
		byte[] encData = f.toBytes();
		try {
			byte[] enc = encrypt(Strings.toBytes(key), encData);
			return f.append('.').appendBase64Url(enc, false).toString();
		} finally {
			f.recycle();
		}
	}

	/** 编码data为jwt token令牌
	 * @param data 编码内容
	 * @param issuer 发布者
	 * @param key 密钥
	 * @param exp 过期时间(分钟为单位)
	 * @return 生成的令牌
	 */
	public static String encode(Map<String, Object> data, String issuer, String key, long exp) throws Exception {
		Date d = new Date(System.currentTimeMillis() + exp * 60_000L);
		return encode(data, issuer, key, d);
	}

	/** 解码jwt token令牌
	 * @param jwtString ACJwt token格式的令牌
	 * @param issuer 发布者
	 * @param key 密钥
	 * @return 解码后的jwt token 的body域
	 */
	public static Map<String, Object> decode(String jwtString, String issuer, String key) throws CodeException {
		if (jwtString == null)
			throw new CodeException(ERR_TOKEN_IS_NULL, "jwtString can't be null");

		List<String> parts = Strings.split(jwtString, ".");
		if (parts.size() != 3 || !Strings.isNotEmpty(parts.get(0), parts.get(1), parts.get(2)))
			throw new CodeException(ERR_BAD_TOKEN, "jwtString error, not {header}.{payload}.{signature} format");

		// 头部校验
		Map<String, Object> headers;
		try {
			headers = base64ToMap(parts.get(0));
		} catch (Exception e) {
			throw new CodeException(ERR_BAD_TOKEN, "jwtString header can't unbase64", e);
		}
		if (headers == null)
			throw new CodeException(ERR_BAD_TOKEN, "jwtString header can't unjson");
		String alg = (String) headers.get("alg");
		if (!HS256.equals(alg))
			throw new CodeException(ERR_NOT_HS256, "jwtString error: alg is not HS256");

		String type = (String) headers.get("typ");
		if (!JWT.equals(type))
			throw new CodeException(ERR_NOT_JWT, "jwtString error: type is not JWT");

		// 解析token的body内容
		Map<String, Object> claims;
		try {
			claims = base64ToMap(parts.get(1));
		} catch (Exception e) {
			throw new CodeException(ERR_BAD_TOKEN, "jwtString claim can't unbase64", e);
		}
		if (claims == null)
			throw new CodeException(ERR_BAD_TOKEN, "jwtString claim can't unjson");

		// 发布者校验
		String iss = (String) claims.get(ISSUER_KEY);
		if (issuer != null && !issuer.equals(iss))
			throw new CodeException(ERR_BAD_ISSUER, "ACJwt error: iss not is right");

		// 校验exp是否过期
		long exp = Convert.asLong(claims.get(EXP_KEY), 0L) * 1000L;
		if (exp < System.currentTimeMillis())
			throw new CodeException(ERR_EXPIRE, "jwtString error: exp is out-dated");

		// 校验签名
		Fmt f = Fmt.get().append(parts.get(0)).append('.').append(parts.get(1));
		byte[] signData = f.toBytes();
		f.recycle();
		try {
			byte[] enc = encrypt(Strings.toBytes(key), signData);
			byte[] sign = Strings.fromBase64Url(parts.get(2));
			if (!Arrays.equals(enc, sign))
				throw new CodeException(ERR_BAD_SIGN, "jwtString error: sign invalid");
		} catch (Exception e) {
			throw new CodeException(ERR_BAD_TOKEN, "jwtString can't calculation sign", e);
		}

		return claims;
	}

	/** 获取jwt token令牌的签名字段文本
	 * @param jwtString 令牌
	 * @return 签名内容
	 */
	public static String getSign(String jwtString) {
		if (jwtString == null || jwtString.isEmpty())
			return "";
		int pos = jwtString.lastIndexOf('.');
		return pos == -1 || pos == jwtString.length() - 1 ? "" : jwtString.substring(pos + 1);
	}

	/** 获取jwt token令牌的claim字段并做json解码
	 * @param jwtString 令牌
	 * @return claim解码后的数据, 解析失败返回null
	 */
	public static Map<String, Object> getClaim(String jwtString) {
		if (jwtString == null || jwtString.isEmpty())
			return null;
		int firstPos = jwtString.indexOf('.');
		int lastPos = jwtString.lastIndexOf('.');
		if (firstPos == -1 || lastPos == -1 || lastPos == jwtString.length() - 1 || firstPos >= lastPos)
			return null;
		return base64ToMap(jwtString.substring(firstPos + 1, lastPos));
	}

	/** 加密 */
	private static byte[] encrypt(byte[] key, byte[] data) throws Exception {
		Mac hmac = (Mac) Mac.getInstance(HMAC_SHA256).clone();
		hmac.init(new SecretKeySpec(key, HMAC_SHA256));
		return hmac.doFinal(data);
	}

	/** base64解码成json对象 */
	private static Map<String, Object> base64ToMap(String base64) {
		byte[] bs = Strings.fromBase64Url(base64);
		return JsonParse.from(new String(bs, StandardCharsets.UTF_8));
	}

}
